// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/plan/plan_builder.h>
#include <pollux/functions/lib/aggregates/tests/utils/AggregationTestBase.h>

using namespace kumo::pollux::exec::test;
using namespace kumo::pollux::functions::aggregate::test;

namespace kumo::pollux::aggregate::test {

namespace {

class MaxSizeForStatsTest : public AggregationTestBase {
 public:
  void SetUp() override {
    AggregationTestBase::SetUp();
  }
};

TEST_F(MaxSizeForStatsTest, nullValues) {
  auto vectors = {make_row_vector({
      make_nullable_flat_vector<int8_t>({std::nullopt, std::nullopt}),
      make_nullable_flat_vector<int16_t>({std::nullopt, std::nullopt}),
      make_nullable_flat_vector<int32_t>({std::nullopt, std::nullopt}),
      make_nullable_flat_vector<int64_t>({std::nullopt, std::nullopt}),
      make_nullable_flat_vector<float>({std::nullopt, std::nullopt}),
      make_nullable_flat_vector<double>({std::nullopt, std::nullopt}),
      make_nullable_flat_vector<bool>({std::nullopt, std::nullopt}),
      make_nullable_flat_vector<int32_t>({std::nullopt, std::nullopt}, DATE()),
      make_nullable_flat_vector<Timestamp>({std::nullopt, std::nullopt}),
      make_nullable_flat_vector<StringView>({std::nullopt, std::nullopt}),
  })};

  testAggregations(
      vectors,
      {},
      {"max_data_size_for_stats(c0)",
       "max_data_size_for_stats(c1)",
       "max_data_size_for_stats(c2)",
       "max_data_size_for_stats(c3)",
       "max_data_size_for_stats(c4)",
       "max_data_size_for_stats(c5)",
       "max_data_size_for_stats(c6)",
       "max_data_size_for_stats(c7)",
       "max_data_size_for_stats(c8)",
       "max_data_size_for_stats(c9)"},
      "SELECT NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL");
}

TEST_F(MaxSizeForStatsTest, nullAndNonNullValues) {
  auto vectors = {make_row_vector({
      make_nullable_flat_vector<int8_t>({std::nullopt, 0}),
      make_nullable_flat_vector<int16_t>({std::nullopt, 0}),
      make_nullable_flat_vector<int32_t>({std::nullopt, 0}),
      make_nullable_flat_vector<int64_t>({std::nullopt, 0}),
      make_nullable_flat_vector<float>({std::nullopt, 0}),
      make_nullable_flat_vector<double>({std::nullopt, 0}),
      make_nullable_flat_vector<bool>({std::nullopt, 0}),
      make_nullable_flat_vector<int32_t>({std::nullopt, 0}, DATE()),
      make_nullable_flat_vector<Timestamp>({std::nullopt, Timestamp(0, 0)}),
      make_nullable_flat_vector<StringView>({std::nullopt, "std::nullopt"}),
  })};

  testAggregations(
      vectors,
      {},
      {"max_data_size_for_stats(c0)",
       "max_data_size_for_stats(c1)",
       "max_data_size_for_stats(c2)",
       "max_data_size_for_stats(c3)",
       "max_data_size_for_stats(c4)",
       "max_data_size_for_stats(c5)",
       "max_data_size_for_stats(c6)",
       "max_data_size_for_stats(c7)",
       "max_data_size_for_stats(c8)",
       "max_data_size_for_stats(c9)"},
      "SELECT 1, 2, 4, 8, 4, 8, 1, 4, 16, 16");
}

template <class T>
T generator(vector_size_t i) {
  return T(i);
}
template <>
Timestamp generator<Timestamp>(vector_size_t i) {
  return Timestamp(i, i);
}
TEST_F(MaxSizeForStatsTest, allScalarTypes) {
  // Make input size at least 8 to ensure drivers get 2 input batches for
  // spilling when tested with TableScan.
  auto vectors = {make_row_vector(
      {make_flat_vector<int64_t>({1, 2, 1, 2, 1, 2, 1, 2}),
       make_flat_vector<int8_t>(8, generator<int8_t>),
       make_flat_vector<int16_t>(8, generator<int16_t>),
       make_flat_vector<int32_t>(8, generator<int32_t>),
       make_flat_vector<int64_t>(8, generator<int64_t>),
       make_flat_vector<float>(8, generator<float>),
       make_flat_vector<double>(8, generator<double>),
       make_flat_vector<bool>(8, generator<bool>),
       make_flat_vector<int32_t>(8, generator<int32_t>, nullptr, DATE()),
       make_flat_vector<Timestamp>(8, generator<Timestamp>)})};

  // With grouping keys.
  testAggregations(
      vectors,
      {"c0"},
      {"max_data_size_for_stats(c1)",
       "max_data_size_for_stats(c2)",
       "max_data_size_for_stats(c3)",
       "max_data_size_for_stats(c4)",
       "max_data_size_for_stats(c5)",
       "max_data_size_for_stats(c6)",
       "max_data_size_for_stats(c7)",
       "max_data_size_for_stats(c8)",
       "max_data_size_for_stats(c9)"},
      "VALUES (1,1,2,4,8,4,8,1,4,16),(2,1,2,4,8,4,8,1,4,16)");

  // Without grouping keys.
  testAggregations(
      vectors,
      {},
      {"max_data_size_for_stats(c1)",
       "max_data_size_for_stats(c2)",
       "max_data_size_for_stats(c3)",
       "max_data_size_for_stats(c4)",
       "max_data_size_for_stats(c5)",
       "max_data_size_for_stats(c6)",
       "max_data_size_for_stats(c7)",
       "max_data_size_for_stats(c8)",
       "max_data_size_for_stats(c9)"},
      "VALUES (1,2,4,8,4,8,1,4,16)");
}

TEST_F(MaxSizeForStatsTest, arrayGlobalAggregate) {
  auto vectors = {make_row_vector({
      make_array_vector<int64_t>({
          {1, 2, 3, 4, 5},
          {},
          {1, 2, 3},
      }),
  })};
  testAggregations(vectors, {}, {"max_data_size_for_stats(c0)"}, "SELECT 44");
}

TEST_F(MaxSizeForStatsTest, mapGlobalAggregate) {
  auto vectors = {make_row_vector({
      make_map_vector<int8_t, int32_t>({
          {{1, 1}, {1, 1}, {1, 1}, {1, 1}, {1, 1}},
          {},
          {{1, 1}, {1, 1}, {1, 1}},
      }),
  })};
  testAggregations(vectors, {}, {"max_data_size_for_stats(c0)"}, "SELECT 29");
}

TEST_F(MaxSizeForStatsTest, rowGlobalAggregate) {
  auto vectors = {make_row_vector({
      make_row_vector({
          make_array_vector<int64_t>({
              {1, 2, 3, 4, 5},
              {},
              {1, 2, 3},
          }),
          make_map_vector<int8_t, int32_t>({
              {{1, 1}, {1, 1}, {1, 1}, {1, 1}, {1, 1}},
              {},
              {{1, 1}, {1, 1}, {1, 1}},
          }),
      }),
  })};
  testAggregations(vectors, {}, {"max_data_size_for_stats(c0)"}, "SELECT 77");
}

TEST_F(MaxSizeForStatsTest, varbinaryGlobalAggregate) {
  VectorPtr varbinaryVector = BaseVector::create(VARBINARY(), 3, pool());
  auto flat_vector = varbinaryVector->as_flat_vector<StringView>();
  flat_vector->set(0, "buf");
  flat_vector->set(1, "");
  flat_vector->set(2, "hello, world !");

  auto vectors = {make_row_vector({varbinaryVector})};
  testAggregations(vectors, {}, {"max_data_size_for_stats(c0)"}, "SELECT 18");
}

TEST_F(MaxSizeForStatsTest, varcharGlobalAggregate) {
  auto vectors = {make_row_vector({
      make_flat_vector<StringView>({
          "{1, 2, 3, 4, 5}",
          "{}",
          "{1, 2, 3}",
      }),
  })};
  testAggregations(vectors, {}, {"max_data_size_for_stats(c0)"}, "SELECT 19");
}

TEST_F(MaxSizeForStatsTest, complexRecursiveGlobalAggregate) {
  auto vectors = {make_row_vector({
      make_row_vector({
          make_flat_vector<StringView>({
              "{1, 2, 3, 4, 5}",
              "{}",
              "{1, 2, 3}",
          }),
          create_map_of_arrays_vector<int8_t, int64_t>({
              {{1, std::nullopt}},
              {{2, {{4, 5, std::nullopt}}}},
              {{3, {{7, 8, 9}}}},
          }),
      }),
  })};

  testAggregations(vectors, {}, {"max_data_size_for_stats(c0)"}, "SELECT 50");
}

TEST_F(MaxSizeForStatsTest, constantEncodingTest) {
  auto columnOne = make_flat_vector<int64_t>({1, 2, 1});
  auto columnTwo = make_row_vector({
      make_flat_vector<StringView>({
          "{1, 2, 3, 4, 5}",
          "{}",
          "{1, 2, 3}",
      }),
      create_map_of_arrays_vector<int8_t, int64_t>({
          {{1, std::nullopt}},
          {{2, {{4, 5, std::nullopt}}}},
          {{std::nullopt, {{7, 8, 9}}}},
      }),
  });
  auto columnTwoConstantEncoded = BaseVector::wrap_in_constant(3, 1, columnTwo);

  auto vectors = {make_row_vector({columnOne, columnTwoConstantEncoded})};

  testAggregations(vectors, {}, {"max_data_size_for_stats(c1)"}, "SELECT 36");

  testAggregations(
      vectors, {"c0"}, {"max_data_size_for_stats(c1)"}, "VALUES (1,36),(2,36)");
}

TEST_F(MaxSizeForStatsTest, dictionaryEncodingTest) {
  auto columnOne = make_flat_vector<int64_t>({1, 2, 1});
  auto columnTwo = make_row_vector({
      make_flat_vector<StringView>({
          "{1, 2, 3, 4, 5}",
          "{}",
          "{1, 2, 3}",
      }),
      create_map_of_arrays_vector<int8_t, int64_t>({
          {{1, std::nullopt}},
          {{2, {{4, 5, std::nullopt}}}},
          {{3, {{7, 8, 9}}}},
      }),
  });
  vector_size_t size = 3;
  auto indices = AlignedBuffer::allocate<vector_size_t>(size, pool());
  auto rawIndices = indices->asMutable<vector_size_t>();
  for (auto i = 0; i < size; ++i) {
    rawIndices[i] = i;
  }
  auto columnTwoDictionaryEncoded =
      BaseVector::wrap_in_dictionary(nullptr, indices, size, columnTwo);
  auto vectors = {make_row_vector({columnOne, columnTwoDictionaryEncoded})};

  testAggregations(vectors, {}, {"max_data_size_for_stats(c1)"}, "SELECT 50");

  testAggregations(
      vectors, {"c0"}, {"max_data_size_for_stats(c1)"}, "VALUES (1,50),(2,36)");
}

TEST_F(MaxSizeForStatsTest, mask) {
  auto data = make_row_vector(
      {"m", "v"},
      {
          make_flat_vector<bool>({false, false, true, true}),
          make_nullable_array_vector<int64_t>({
              std::nullopt,
              std::nullopt,
              {{1, 2, 3, 4}},
              {{1, 2, 3, 4, 5}},
          }),
      });

  auto plan = PlanBuilder()
                  .values({data})
                  .singleAggregation({}, {"max_data_size_for_stats(v)"}, {"m"})
                  .planNode();
  assertQuery(plan, "SELECT 44");
}

} // namespace
} // namespace kumo::pollux::aggregate::test
